from __future__ import annotations

import os, json
import logging
import re
from pathlib import Path

from typing import Any, Dict, Optional

from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters

from mcp.client.stdio import stdio_client


logger = logging.getLogger(__name__)


def _setup_audit_logger() -> logging.Logger:
    """创建专用 MCP 审计日志记录器"""
    log_path = Path(os.getenv("MCP_LOG_PATH", "logs/mcp.log"))
    try:
        log_path.parent.mkdir(parents=True, exist_ok=True)
    except Exception as e:
        # 如果目录创建失败，退回到当前目录
        logger.warning(f"无法创建 MCP 日志目录 {log_path.parent}: {e}")
        log_path = Path("mcp.log")

    audit_logger = logging.getLogger("mcp.audit")
    if not audit_logger.handlers:
        handler = logging.FileHandler(log_path, encoding="utf-8")
        handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
        audit_logger.addHandler(handler)
        audit_logger.setLevel(logging.INFO)
        audit_logger.propagate = False
    return audit_logger


_mcp_audit_logger = _setup_audit_logger()


def _serialize_for_log(obj: Any) -> str:
    """将对象转换为日志友好的 JSON 字符串"""
    try:
        return json.dumps(obj, ensure_ascii=False, default=str)
    except Exception:
        return str(obj)


def _read_mcp_args_from_env_file() -> Optional[str]:
    """从 .env 文件直接读取 MCP_ARGS_JSON（支持多行格式）"""
    env_file = Path(".env")
    if not env_file.exists():
        return None
    
    try:
        lines = env_file.read_text(encoding="utf-8").splitlines()
        collecting = False
        json_lines = []
        
        for line in lines:
            stripped = line.strip()
            
            # 找到 MCP_ARGS_JSON 开始
            if stripped.startswith("MCP_ARGS_JSON="):
                # 提取等号后面的部分
                value_part = stripped.split("=", 1)[1].strip()
                if value_part:
                    json_lines.append(value_part)
                # 检查是否已结束（闭合括号）
                if value_part.endswith("]"):
                    break
                collecting = True
                continue
            
            # 如果正在收集多行值
            if collecting:
                # 移除行尾注释
                stripped = re.sub(r'\s*#.*$', '', stripped)
                if stripped:
                    json_lines.append(stripped)
                # 如果遇到闭合括号，结束收集
                if stripped.endswith("]"):
                    break
                # 如果遇到新的变量定义（以大写字母开头），结束收集
                if re.match(r'^[A-Z_][A-Z0-9_]*=', stripped):
                    break
        
        if json_lines:
            # 合并所有行，移除多余空格
            value = ' '.join(json_lines)
            # 展开环境变量引用 ${VAR} 或 $VAR
            value = re.sub(r'\$\{([^}]+)\}', lambda m: os.getenv(m.group(1), ''), value)
            value = re.sub(r'\$([A-Z_][A-Z0-9_]*)', lambda m: os.getenv(m.group(1), ''), value)
            return value
            
    except Exception as e:
        logger.debug(f"Failed to read MCP_ARGS_JSON from .env file: {e}")
    
    return None


class MCPManager:

    def __init__(self) -> None:

        self.session: Optional[ClientSession] = None

        self.exit_stack = AsyncExitStack()

        self.stdio = None

        self.write = None
        self.enabled = False  # 标记是否成功启用


    async def start(self) -> None:

        if self.session:

            return  # already started

        command = os.getenv("MCP_COMMAND", "docker")

        args_json = os.getenv("MCP_ARGS_JSON", "[]")
        
        # 如果环境变量值太短（可能是多行被截断），尝试从 .env 文件读取
        if len(args_json) <= 2 and (args_json == "[]" or args_json == "["):
            logger.debug("MCP_ARGS_JSON from env seems truncated, trying to read from .env file")
            env_file_value = _read_mcp_args_from_env_file()
            if env_file_value:
                args_json = env_file_value
                logger.debug(f"Read MCP_ARGS_JSON from .env file: {args_json[:100]}...")

        if not args_json or args_json == "[]":

            logger.warning("MCP_ARGS_JSON is empty or not set, MCP will be disabled")

            self.enabled = False

            return

        try:

            args = json.loads(args_json)

        except json.JSONDecodeError as e:

            logger.error(f"Failed to parse MCP_ARGS_JSON: {e}, MCP will be disabled")

            self.enabled = False

            return

        if not args:

            logger.warning("MCP_ARGS_JSON is empty array, MCP will be disabled")

            self.enabled = False

            return

        try:

            params = StdioServerParameters(command=command, args=args, env=None)

            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(params))

            self.stdio, self.write = stdio_transport

            self.session = await self.exit_stack.enter_async_context(

                ClientSession(self.stdio, self.write)

            )

            await self.session.initialize()  # handshake 完成

            self.enabled = True

            logger.info("MCP session started successfully")

        except Exception as e:

            logger.error(f"Failed to start MCP session: {e}, MCP will be disabled")

            self.enabled = False

            # 确保清理资源

            try:

                await self.exit_stack.aclose()

            except:

                pass

            self.session = None

            self.stdio = None

            self.write = None



    async def stop(self) -> None:

        await self.exit_stack.aclose()

        self.session = None

        self.stdio = None

        self.write = None



    async def list_tools(self):

        if not self.session or not self.enabled:

            raise RuntimeError("MCP not started or disabled")

        res = await self.session.list_tools()

        return res.tools  # MCP 工具列表（含 JSON Schema）



    async def call_tool(self, name: str, args: Dict[str, Any]):

        if not self.session or not self.enabled:

            raise RuntimeError("MCP not started or disabled")

        _mcp_audit_logger.info(
            "CALL %s %s",
            name,
            _serialize_for_log(args),
        )

        try:
            result = await self.session.call_tool(name, args)
            _mcp_audit_logger.info(
                "RESULT %s %s",
                name,
                _serialize_for_log(result),
            )
            return result
        except Exception as e:
            _mcp_audit_logger.exception(
                "ERROR %s %s -> %s",
                name,
                _serialize_for_log(args),
                e,
            )
            raise



mcp = MCPManager()

